#include "Encoder.h"
#include <fstream>
#include <stdlib.h>
#include "Decoder.h"
#include "NsightDebug.h"

using namespace std;

__global__ void MEStep2Kernel(BlockList currentFrame, BlockList previousFrame, double* data, cudaMADInfoContainer* result, int columnLength, int meRange)
{
	int matrixSize = columnLength * columnLength;
	int meBlocksPerRow = 2 * meRange + 1;
	int mePoolSize = meBlocksPerRow * meBlocksPerRow;

	int targetBlockX = blockIdx.x * columnLength;
	int targetBlockY = blockIdx.y * columnLength;
	
	int blockIndex = blockIdx.y * gridDim.x + blockIdx.x;
	int meIndex = threadIdx.y * blockDim.x + threadIdx.x;

	int vecX = threadIdx.x - meRange;
	int vecY = threadIdx.y - meRange;

	int meCol = targetBlockX + vecX;
	int meRow = targetBlockY + vecY;

	extern __shared__ double s[];

	s[meIndex+2] = 0;

	bool validBlock = true;

	for(int i = 0; i < matrixSize && validBlock; ++i)
	{
		int val = data[blockIndex * matrixSize * mePoolSize + meIndex * matrixSize + i];
		if(val == INT_MAX)
		{
			validBlock = false;
			s[meIndex+2] = INT_MAX;
		}
		else
			s[meIndex+2] += fabsf(val);
	}

	s[meIndex+2] /= matrixSize;

	__syncthreads();

	s[1] = INT_MAX;

	if(meIndex == 0)
	{
		for(int i = 0; i < mePoolSize; ++i)
		{
			if(s[i+2] < s[1])
			{
				s[1] = s[i+2];
				s[0] = i;

			}
		}
	}

	__syncthreads();

	if(meIndex == (int)s[0])
	{
		cudaMADInfoContainer returnValue;
		returnValue.MAD = s[1];
		returnValue.x = vecX;
		returnValue.y = vecY;
		returnValue.targetX = targetBlockX;
		returnValue.targetY = targetBlockY;
		returnValue.type = returnValue.MAD == 0 ? 0 : 2;

		result[blockIndex] = returnValue;
	}
}

__global__ void MEStep1Kernel(BlockList previousFrame, BlockList currentFrame, int meRange, double* d, int columnLength)
{
	int matrixSize = currentFrame.getMatrixSize();
	int targetBlockX = blockIdx.x * currentFrame.getColumnLength();
	int targetBlockY = blockIdx.y * currentFrame.getColumnLength();
	
	int blockIndex = blockIdx.y * gridDim.x + blockIdx.x;
	int MADIndex = blockIdx.z;
	int valueIndex = threadIdx.y * blockDim.x + threadIdx.x;

	int meBlocksPerRow = 2 * meRange + 1;
	int mePoolSize = meBlocksPerRow * meBlocksPerRow;

	int meX = MADIndex % meBlocksPerRow - meRange;
	int meY = MADIndex / meBlocksPerRow - meRange;

	int meCol = targetBlockX + meX + threadIdx.x;
	int meRow = targetBlockY + meY + threadIdx.y;

	bool outOfBounds = meRow < 0 || meRow > currentFrame.getDataHeight()-columnLength || meCol < 0 || meCol > currentFrame.getDataWidth()-columnLength;

	if(!outOfBounds)
	{
		double targetPixel = currentFrame.getPixel(targetBlockY + threadIdx.y, targetBlockX + threadIdx.x);
		double comparingPixel = previousFrame.getPixel(meRow, meCol);
		d[blockIndex * matrixSize * mePoolSize + MADIndex * matrixSize + valueIndex] = targetPixel - comparingPixel;
	}
	else
		d[blockIndex * matrixSize * mePoolSize + MADIndex * matrixSize + valueIndex] = INT_MAX;
}

cudaMADInfoContainer* cudaCalculateME(BlockList& previousFrame, BlockList& currentFrame, int meRange = 2, int columnLength = 16)
{
	int meBlocksPerRow = 2 * meRange + 1;

	double* d;

	int totBlocks = currentFrame.getRows() * currentFrame.getColumns();

	cudaMADInfoContainer* cuMADInfo;
	cudaMADInfoContainer* MADInfo = new cudaMADInfoContainer[totBlocks];

	cudaMalloc(&cuMADInfo, totBlocks * sizeof(cudaMADInfoContainer));

	cudaMalloc(&d, totBlocks * meBlocksPerRow*meBlocksPerRow * currentFrame.getMatrixSize() * sizeof(double));

	uint8_t* cuPreviousFrameData;
	uint8_t* cuCurrentFrameData;

	previousFrame.cloneToDevice(&cuPreviousFrameData);
	currentFrame.cloneToDevice(&cuCurrentFrameData);

	BlockList cuPreviousFrame;
	BlockList cuCurrentFrame;

	cuPreviousFrame.initialize(cuPreviousFrameData, previousFrame.getDataWidth(), previousFrame.getDataHeight(), totBlocks, previousFrame.getColumnLength());
	cuCurrentFrame.initialize(cuCurrentFrameData, currentFrame.getDataWidth(), currentFrame.getDataHeight(), totBlocks, currentFrame.getColumnLength());

	dim3 blocks(currentFrame.getColumns(), currentFrame.getRows(), meBlocksPerRow*meBlocksPerRow);
	dim3 threadsPerBlock(columnLength, columnLength);

	MEStep1Kernel<<<blocks, threadsPerBlock>>>(cuPreviousFrame, cuCurrentFrame, meRange, d, columnLength);

	cudaDeviceSynchronize();

	dim3 blocks2(currentFrame.getColumns(), currentFrame.getRows());
	dim3 threadsPerBlock2(meBlocksPerRow, meBlocksPerRow);

	int sharedSize = (meBlocksPerRow*meBlocksPerRow + 2) * sizeof(double);

	MEStep2Kernel<<<blocks2,threadsPerBlock2, sharedSize>>>(cuCurrentFrame, cuPreviousFrame, d, cuMADInfo, columnLength, meRange);

	cudaMemcpy(MADInfo, cuMADInfo, totBlocks * sizeof(cudaMADInfoContainer), cudaMemcpyDeviceToHost);
	
	cudaFree(cuMADInfo);
	cudaFree(d);

	return MADInfo;
}

__global__ void IFrameKernel(BlockList data, int* qData, int columnLength)
{
	Normalization n;
	Quantization q(qData, columnLength);
	DCT dct(columnLength);

	int index = blockIdx.x * blockDim.x + threadIdx.x;
	Matrix m = data[index];
	n.performNormalization(m);
	dct.performFDCT(m);
	q.performQuantization(m);
	data.setDevice(index, m);

}

void cudaEncodeIFrame(Normalization& n, DCT& dct, Quantization& q, BlockList& data)
{
	int blocks = data.getRows();
	int threadsPerBlock = data.getColumns();

	int* quantData = q.copyToDevice();
	int columnLength = q.getQuantizationMatrix().getColumnLength();

	IFrameKernel<<<blocks, threadsPerBlock>>>(data, quantData, columnLength);

	cudaFree(quantData);
}

Encoder::Encoder(QuantizationMatrix* qm, bool rle, int width, int height,  int matrixColumnSize)
{
	/*Size*/
	_matrixColumnLength = matrixColumnSize;
	_matrixSize = matrixColumnSize*matrixColumnSize;

	/*Vars*/
	_RLE = rle;
	_width = width;
	_height = height;
	_columnSize = matrixColumnSize;

	/*BlockList*/
	_totalBlocks = (_width/matrixColumnSize)*(_height/matrixColumnSize);

	/*Modules*/
	_rleModule = new RLE(_totalBlocks);
	_normModule = new Normalization();
	_dctModule = new DCT(_matrixColumnLength);
	_quantModule = new Quantization(qm);
	_quantizationMatrix = qm;
	_meModule = new ME(_columnSize, _width, _height, _meRange);
}

Encoder::Encoder(EncoderSettings &config, QuantizationMatrix* qm, int matrixColumnSize)
{
	/*Size*/
	_matrixColumnLength = matrixColumnSize;
	_matrixSize = matrixColumnSize*matrixColumnSize;


	/*Vars*/
	_RLE = config.getRLE();
	_width = config.getWidth();
	_height = config.getHeight();
	_columnSize = matrixColumnSize;
	_meRange = config.getMerange();

	/*BlockList*/
	_totalBlocks = (_width/matrixColumnSize)*(_height/matrixColumnSize);

	/*Modules*/
	_rleModule = new RLE(_totalBlocks);
	_normModule = new Normalization();
	_dctModule = new DCT(_matrixColumnLength);
	_quantModule = new Quantization(qm);
	_quantizationMatrix = qm;
	_meModule = new ME(_columnSize, _width, _height, _meRange);

	/*File names*/
	_logFile = config.logfile();
}

util::BitStreamWriter* Encoder::encode(uint8_t* buff, bool useCUDA)
{
	NsightDebug::pushTracerStack("Norm-FDCT-Quant-RLE", 0xffff0000);
	_blockList.initialize(buff, _width, _height, _totalBlocks, _columnSize);

	int size = _totalBlocks*_matrixSize;

	if(useCUDA)
	{
		uint8_t* CUData;

		_blockList.cloneToDevice(&CUData);

		BlockList CUList;

		CUList.initialize(CUData, _width, _height, _totalBlocks, _columnSize);

		cudaEncodeIFrame(*_normModule, *_dctModule, *_quantModule, CUList);

		_blockList.cloneToHost(CUList.data());
	}
	else
	{
#if I_ENC_PARALLEL == 1
#pragma omp parallel for
#endif
		for(int i = 0; i < _blockList.size(); ++i)
		{
			Matrix m = _blockList[i];
			_normModule->performNormalization(m);
			_dctModule->performFDCT(m);
			_quantModule->performQuantization(m);
			_blockList.set(i, m);
		}
	}
	int8_t* zigzag;
	zigzag = _rleModule->performZigZag(_blockList);

	uint8_t* buffer = new uint8_t[size];
	util::BitStreamWriter* bsw = new util::BitStreamWriter(buffer, size);

	if(_RLE)
		_rleModule->performRLE(zigzag, bsw);
	else
		_rleModule->performRawEncoding(zigzag, bsw);

	delete[] zigzag;

	if(bsw->get_position()%8 != 0)
		bsw->flush();

	_blockList.clear();

	NsightDebug::pop();

	return bsw;
}

util::BitStreamWriter* Encoder::encodePFrame(uint8_t* orig, uint8_t *buff, uint8_t* previousFrame, bool useCUDA)
{
	/*Initialize list of macroblocks*/
	_blockList.initialize(buff, _width, _height, _totalBlocks, _columnSize);
	BlockList sourceFrame;

	sourceFrame.initialize(orig, _width, _height, _totalBlocks, _columnSize);

	BlockList blpf;
	blpf.initialize(previousFrame, _width, _height, _totalBlocks, _columnSize);

	int size = _totalBlocks*_matrixSize;

	uint8_t* buffer = new uint8_t[size];

	util::BitStreamWriter* bsw = new util::BitStreamWriter(buffer, size);

	util::BitStreamWriter** bswArray = new util::BitStreamWriter*[_totalBlocks];

	if(useCUDA)
	{
		NsightDebug::pushTracerStack("ME-MC", 0xffffff00);
		cudaMADInfoContainer* MADInfo = cudaCalculateME(blpf, _blockList, _meRange, _columnSize);
		NsightDebug::pop();

		NsightDebug::pushTracerStack("PFrame Data Encoding", 0xff00ffff);
		for(int i = 0; i < _blockList.size(); ++i)
		{
			cudaMADInfoContainer MAD = MADInfo[i];
			MADInfoContainer temp;
			temp.MAD = MAD.MAD;
			temp.type = MAD.type;
			temp.x = MAD.x;
			temp.y = MAD.y;
			Matrix comparingBlock = blpf.getBlockAtCoord(MAD.targetX + MAD.x, MAD.targetY + MAD.y);
			temp.comparingBlock = comparingBlock;
			temp.diff = _blockList.getBlockAtCoord(MAD.targetX, MAD.targetY) - comparingBlock;
			compressPFrame(sourceFrame, bswArray, temp, i);
		}

		delete[] MADInfo;
		NsightDebug::pop();
	}
	else
	{
#if P_ENC_PARALLEL == 1
#pragma omp parallel for
#endif
		for(int i = 0; i < _blockList.size(); ++i)
		{
			Matrix m = _blockList[i];

			int row = i / (_width / _columnSize) * _columnSize;
			int col = i % (_width / _columnSize) * _columnSize;

			NsightDebug::pushTracerStack("ME-MC", 0xffffff00);

			MADInfoContainer MADInfo = _meModule->calculateMotionVector(m, col, row, blpf);

			NsightDebug::pop();

			NsightDebug::pushTracerStack("PFrame Data Encoding", 0xff00ffff);

			compressPFrame(sourceFrame, bswArray, MADInfo, i);

			NsightDebug::pop();
		}
	}

	for(int k = 0; k < _totalBlocks; ++k)
	{
		util::BitStreamWriter* temp = bswArray[k];

util:BitStreamReader bsr(temp->get_buffer(), temp->get_size());
		for(int j = 0; j < temp->get_position(); ++j)
			bsw->put_bit(bsr.get_bit());

		delete temp;
	}

	delete[] bswArray;

	_blockList.clear();

	return bsw;
}

Encoder::~Encoder()
{
	delete _rleModule;
	delete _quantModule;
	delete _normModule;
	delete _dctModule;
	delete _meModule;
}

void Encoder::compressPFrame(BlockList &sourceFrame, util::BitStreamWriter** bswArray, MADInfoContainer &MADInfo, int index)
{
	util::BitStreamWriter* blockWriter = new util::BitStreamWriter(_matrixSize);
	bswArray[index] = blockWriter;

	blockWriter->put(2, MADInfo.type);
	if(MADInfo.type == 0)
	{
		if(MADInfo.x < 0)
			blockWriter->put_bit(1);
		else
			blockWriter->put_bit(0);
		blockWriter->put(7, (uint32_t)MADInfo.x);
		if(MADInfo.y < 0)
			blockWriter->put_bit(1);
		else
			blockWriter->put_bit(0);
		blockWriter->put(7, (uint32_t)MADInfo.y);
	}
	else if(MADInfo.type == 1)
	{
		blockWriter->put(8, (uint32_t)MADInfo.x);
		blockWriter->put(8, (uint32_t)MADInfo.y);
		Encoder matrixEncoder(_quantizationMatrix, _RLE, 16, 16, 4);
		uint8_t* macroBlockData = MADInfo.diff.getMatrixData();
		util::BitStreamWriter* encodedMacroBlock = matrixEncoder.encode(macroBlockData, false);
		blockWriter->put_bit(0);
		uint32_t bytesWritten = (uint32_t)(encodedMacroBlock->get_position()/8);
		if(encodedMacroBlock->get_position()%8 != 0)
			++bytesWritten;
		blockWriter->put(31, bytesWritten);
		util::BitStreamReader encodedMacroBlockReader(encodedMacroBlock->get_buffer(), encodedMacroBlock->get_size());
		for(int k = 0; k < encodedMacroBlock->get_position(); ++k)
			blockWriter->put_bit(encodedMacroBlockReader.get_bit());

		encodedMacroBlock->clear();
		delete encodedMacroBlock;
		delete[] macroBlockData;
	}
	else
	{
		Encoder matrixEncoder(_quantizationMatrix, _RLE, 16, 16, 4);
		//uint8_t* macroBlockData = (MADInfo.comparingBlock).getMatrixData();
		uint8_t* macroBlockData = sourceFrame[index].getMatrixData();
		//uint8_t* macroBlockData = MADInfo.diff.getMatrixData();
		util::BitStreamWriter* encodedMacroBlock  = matrixEncoder.encode(macroBlockData, false);
		blockWriter->put_bit(0);
		uint32_t bytesWritten = (uint32_t)(encodedMacroBlock->get_position()/8);
		if(encodedMacroBlock->get_position()%8 != 0)
			++bytesWritten;
		blockWriter->put(31, bytesWritten);
		util::BitStreamReader encodedMacroBlockReader(encodedMacroBlock->get_buffer(), encodedMacroBlock->get_size());
		for(int k = 0; k < encodedMacroBlock->get_position(); ++k)
			blockWriter->put_bit(encodedMacroBlockReader.get_bit());

		encodedMacroBlock->clear();
		delete encodedMacroBlock;
		delete[] macroBlockData;
	}
}
